#!/usr/bin/env python
# -*- coding: utf-8; py-indent-offset:4 -*-

"""
indicators.py -- 公式指标库（参考通达信、文华、同花顺等）
------------------------------------------ 
@Author:    WakeBlade
@Email:     2390245@qq.com
@Date:      2022-10-7
------------------------------------------
"""

from typing import Union, List, Sequence, Any
import pandas as pd
import numpy as np

#####################################################################
# 0.公共函数
#####################################################################
def _asArray(s:Union[List, Sequence, np.ndarray, pd.Series]):
    if isinstance(s, np.ndarray):
        return np.copy(s)
    if isinstance(s, pd.Series):
        return s.values
    return np.array(s)

def _asSeries(s:Union[List, Sequence, np.ndarray, pd.Series]):
    return s.copy() if isinstance(s, pd.Series) else pd.Series(s)

def _ensure(s:Union[List, Sequence, np.ndarray, pd.Series], asArray:bool = False):
    return _asArray(s) if asArray else _asSeries(s)

def _col(df:pd.DataFrame, c:str, asArray:bool = False):
    if df is None or c is None or len(c)<1:
        return None

    if c in df:
        s = df[c]
    elif c.capitalize() in df:
        s = df[c.capitalize()]
    elif c.upper() in df:
        s = df[c.uppser()]
    elif c.lower() in df:
        s = df[c.lower()]
    else:
        return None
    
    return _ensure(s, asArray)

#####################################################################
# 1.基础行情指标
#####################################################################
def OPEN(bars:pd.DataFrame, asArray:bool = False):
    col =  _col(bars, "open", asArray) 
    return col if col is not None else _col(bars, "open_price", asArray)

def HIGH(bars:pd.DataFrame, asArray:bool = False):
    col =  _col(bars, "high", asArray) 
    return col if col is not None else _col(bars, "high_price", asArray)

def LOW(bars:pd.DataFrame, asArray:bool = False):
    col =  _col(bars, "low", asArray) 
    return col if col is not None else _col(bars, "low_price", asArray)

def CLOSE(bars:pd.DataFrame, asArray:bool = False):
    col =  _col(bars, "close", asArray) 
    return col if col is not None else _col(bars, "close_price", asArray)

def VOLUME(bars:pd.DataFrame, asArray:bool = False):
    return _col(bars, "volume", asArray) 

def TURNOVER(bars:pd.DataFrame, asArray:bool = False):
    return _col(bars, "turnover", asArray) 

def OPEN_INTEREST(bars:pd.DataFrame, asArray:bool = False):
    return _col(bars, "open_interest", asArray) 

#定义函数的缩略版
# O, H, L, C, V, T, I = OPEN, HIGH, LOW, CLOSE, VOLUME, TURNOVER, OPEN_INTEREST

#####################################################################
## 2.均线系统
#####################################################################
# 算数平均
def MA(s:Union[List, Sequence, np.ndarray, pd.Series], n:int, min_periods:int = 1, asArray:bool = False):
    ma = _asSeries(s).rolling(n, min_periods=min_periods).mean()
    return _ensure(ma, asArray)

# 前向算数平均(未来函数)
def XMA(s:Union[List, Sequence, np.ndarray, pd.Series], n:int, asArray:bool = False):
    r = _ensure(s, asArray)
    for i in range(len(r)):
        if i+n>len(r):
            r[i] = np.mean(r[i:])
        else:
            r[i] = np.mean(r[i:i+n])
    return r

# 指数平滑移动平均n
def EMA(s:Union[List, Sequence, np.ndarray, pd.Series], n:int, min_periods:int = 1, asArray:bool = False):
    ema = _asSeries(s).ewm(span=n, adjust=False, min_periods=min_periods).mean()
    return _ensure(ema, asArray)

# 简单移动平均
def SMA(s:Union[List, Sequence, np.ndarray, pd.Series], n:int, m:int=1, min_periods:int = 1, asArray:bool = False):
    ma = _asSeries(s).rolling(n, min_periods=min_periods).mean()
    for i in range(n+1, len(ma)):
        ma.iloc[i] = (m * s[i] + (n - m) * ma.iloc[i-1]) / n  
    return _asArray(ma, asArray)

# 中位数
def MEDIAN(s:Union[List, Sequence, np.ndarray, pd.Series], n:int, min_periods:int = 1, asArray:bool = False):
    r = _asSeries(s).rolling(n, min_periods=min_periods).median()
    return _ensure(r, asArray)

# 分位数
def QUANTILE(s:Union[List, Sequence, np.ndarray, pd.Series], n:int, q:int = 50, min_periods:int = 1, asArray:bool = False):
    r = _asSeries(s).rolling(n, min_periods=min_periods).quantile(q)
    return _ensure(r, asArray)

# 标准差
def STD(s:Union[List, Sequence, np.ndarray, pd.Series], n:int, min_periods:int = 1, asArray:bool = False):
    r = _asSeries(s).rolling(n, min_periods=min_periods).std()
    return _ensure(r, asArray)

#####################################################################
## 3.算数计算
#####################################################################
# 四舍五入
def ROUND(s:Union[List, Sequence, np.ndarray, pd.Series], n:int = 2, asArray:bool = False):
    return np.round(_ensure(s, asArray), n)

# 向上取整
def CEIL(s:Union[List, Sequence, np.ndarray, pd.Series], asArray:bool = False):
    return np.ceil(_ensure(s, asArray))

# 向下取整
def FLOOR(s:Union[List, Sequence, np.ndarray, pd.Series], asArray:bool = False):
    return np.floor(_ensure(s, asArray))

# 绝对值
def ABS(s:Union[List, Sequence], asArray:bool = False):
    return np.abs(_ensure(s, asArray))

# 取符号，>0变1，<0变-1
def SIGN(s:Union[List, Sequence, np.ndarray, pd.Series], asArray:bool = False):
    return np.sign(_ensure(s, asArray))

# 取最大值
def MAX(s1:Union[List, Sequence, np.ndarray, pd.Series], s2:Union[List, Sequence, np.ndarray, pd.Series], asArray:bool = False):
    r = np.maximum(_ensure(s1, asArray), _ensure(s2, asArray))
    return r

# 取最小值
def MIN(s1:Union[List, Sequence, np.ndarray, pd.Series], s2:Union[List, Sequence, np.ndarray, pd.Series], asArray:bool = False):
    r = np.minimum(_ensure(s1, asArray), _ensure(s2, asArray))
    return r

#####################################################################
## 4.逻辑运算
#####################################################################
# IF三元运算
def IF(x:Union[List, Sequence, np.ndarray, pd.Series], a:Union[List, Sequence, np.ndarray, pd.Series], b:Union[List, Sequence, np.ndarray, pd.Series], asArray:bool = False):
    # assert not(x is None or a is None or b is None), "IF三元运算，x,a,b不能为空！" 
    # assert len(x)==len(a) and len(x)==len(b), "IF三元运算，x,a,b必须同长度！"
    _x = _asSeries(x)
    _a = a if isinstance(a,int) or isinstance(a,float) else _asSeries(a)
    _b = b if isinstance(b,int) or isinstance(b,float) else _asSeries(b)
    return _ensure(_x * _a + (1-_x) * _b, asArray)

# AND运算
def AND(a:Union[List, Sequence, np.ndarray, pd.Series], b:Union[List, Sequence, np.ndarray, pd.Series], asArray:bool = False):
    assert not(a is None or b is None), "AND二元运算，a,b不能为空！" 
    assert len(a)==len(b), "AND二元运算，a,b必须同长度！"
    _a, _b = _ensure(a, asArray), _ensure(b, asArray)
    return np.logical_and(_a, _b)

# OR运算
def OR(a:Union[List, Sequence, np.ndarray, pd.Series], b:Union[List, Sequence, np.ndarray, pd.Series], asArray:bool = False):
    assert not(a is None or b is None), "OR二元运算，a,b不能为空！" 
    assert len(a)==len(b), "OR二元运算，a,b必须同长度！"
    _a, _b = _ensure(a, asArray), _ensure(b, asArray)
    return np.logical_or(_a, _b)

# XOR运算
def XOR(a:Union[List, Sequence, np.ndarray, pd.Series], b:Union[List, Sequence, np.ndarray, pd.Series], asArray:bool = False):
    assert not(a is None or b is None), "XOR二元运算，a,b不能为空！" 
    assert len(a)==len(b), "XOR二元运算，a,b必须同长度！"
    _a, _b = _ensure(a, asArray), _ensure(b, asArray)
    return np.logical_xor(_a, _b)

# NOT运算
def NOT(a:Union[List, Sequence, np.ndarray, pd.Series], asArray:bool = False):
    if a is None:
        return None
    return _ensure(a==False, asArray)

# 红绿柱
def ISRED(bars:pd.DataFrame, asArray:bool = False):
    return CLOSE(bars, asArray) > OPEN(bars, asArray)

# 是否上涨
def ISUP(s:Union[List, Sequence, np.ndarray, pd.Series], asArray:bool = False):

    return _ensure(_asSeries(s).astype(int).diff()>0, asArray)

#####################################################################
## 5.引用
#####################################################################
# 取指定位置值
def get(s:Union[List, Sequence, np.ndarray, pd.Series], n:int = -1):
    return _asArray(s)[n]

# 偏移n
def REF(s:Union[List, Sequence, np.ndarray, pd.Series], n:Union[List, Sequence, np.ndarray, pd.Series, int] = 1, fill_value:Any = None, asArray:bool = False):
    r = _asSeries(s)
    if isinstance(n,int):
        r = r.shift(n, fill_value=fill_value)
    else:
        i = r.index.argsort()
        r = r[(i-n)*(i>=n)]
    return _ensure(r, asArray)

# 逆序
def REVERSE(s:Union[List, Sequence, np.ndarray, pd.Series], asArray:bool = False):
    r = _ensure(s, asArray)
    return s[::-1]

# 求日差
def DIFF(s:Union[List, Sequence, np.ndarray, pd.Series], asArray:bool = False):
    r = _asSeries(s).astype(float).diff()
    return _ensure(r, asArray)

# 求日比
def PCT(s:Union[List, Sequence, np.ndarray, pd.Series], asArray:bool = False):
    r = _asSeries(s).pct_change()
    return _ensure(r, asArray)

# 取区间内最大值
def HHV(s:Union[List, Sequence, np.ndarray, pd.Series], n:int = 3, min_periods:int = 1, asArray:bool = False):
    r = _asSeries(s).rolling(n, min_periods=min_periods).max()
    return _ensure(r, asArray)

# 取区间内最大值离当前周期数
def HHVBARS(s:Union[List, Sequence, np.ndarray, pd.Series], n:int = 3, min_periods:int = 1, asArray:bool = False):
    r = _asSeries(s).rolling(n, min_periods=min_periods).apply(
        lambda x: len(x)-((x.index.argsort()+1)*(x==x.max())).max()
    )
    return _ensure(r, asArray)

# 取区间内最小值
def LLV(s:Union[List, Sequence, np.ndarray, pd.Series], n:int = 3, min_periods:int = 1, asArray:bool = False):
    r = _asSeries(s).rolling(n, min_periods=min_periods).min()
    return _ensure(r, asArray)

# 取区间内最小值离当前周期数
def LLVBARS(s:Union[List, Sequence, np.ndarray, pd.Series], n:int = 3, min_periods:int = 1, asArray:bool = False):
    r = _asSeries(s).rolling(n, min_periods=min_periods).apply(
        lambda x: len(x)-((x.index.argsort()+1)*(x==x.min())).max()
    )
    return _ensure(r, asArray)

# 区间求和
def SUM(s:Union[List, Sequence, np.ndarray, pd.Series], n:int = 3, min_periods:int = 1, asArray:bool = False):
    r = _asSeries(s).rolling(n, min_periods=min_periods).sum()
    return _ensure(r, asArray)

# 区间求积
def PROD(s:Union[List, Sequence, np.ndarray, pd.Series], n:int = 3, asArray:bool = False):
    r = _ensure(s, asArray)
    for i in range(len(r), 1, -1):
        if i<n:
            r[i-1] = np.prod(r[:i])
        else:
            r[i-1] = np.prod(r[i-n:i])
    return r

# 求N日内满足条件的次数
def COUNT(s:Union[List, Sequence, np.ndarray, pd.Series], n:int = 3, min_periods:int = 1, asArray:bool = False):
    assert s is not None and isinstance(s[0], bool), "s必须非空，且为bool类型"
    r = _asSeries(s).rolling(n, min_periods=min_periods).sum()
    return _ensure(r, asArray)

# N天内都出现
def EVERY(s:Union[List, Sequence, np.ndarray, pd.Series], n:int = 3, asArray:bool = False):
    return COUNT(s, n, asArray)==n  

# N天到m天内持续出现
def LAST(s:Union[List, Sequence, np.ndarray, pd.Series], n:int = 5, m:int = 3, asArray:bool = False):                
    assert n>=m, "计算LAST(s, n, m)，n不能小于m！"
    return (COUNT(s, n, asArray)-COUNT(s, m, asArray))==(n-m)

# N天内至少出现过一次
def EXIST(s:Union[List, Sequence, np.ndarray, pd.Series], n:int = 3, asArray:bool = False):
    return COUNT(s, n, asArray)>0  

# N天内未出现过一次
def NOEXIST(s:Union[List, Sequence, np.ndarray, pd.Series], n:int = 3, asArray:bool = False):
    """
    NOEXIST: N天内未出现过一次
    -----------------------------
    Params:
        s - 条件判断序列
        n - n周期内
        asArray - 是否转为np.array
    
    Return:
        n周期内满足条件次数为0
    """

    return COUNT(s, n, asArray)==0  

# 有效周期数
def BARSCOUNT(s:Union[List, Sequence, np.ndarray, pd.Series], asArray:bool = False):
    r = _asSeries(s)
    return _ensure(r.index.argsort()+1, asArray)

# 上一次出现至今周期数(-1代表无出现，0代表当前满足条件，N代表N日前满足条件)
def BARSLAST(s:Union[List, Sequence, np.ndarray, pd.Series], asArray:bool = False):
    """
    BARSLAST: 最近一次满足条件离当前多少周期
    -----------------------------
    Params:
        s - 条件判断序列
        asArray - 是否转为np.array
    
    Return:
        最近满足条件离当前k周期数
    """

    r = _asSeries(s)
    i = r.index.argsort()+1
    p = (i * r).cummax()
    return _ensure((p>0) * (i-p) - (p<1), asArray)

# (未来函数）下一次出现至今周期数(-1代表无出现，0代表当前满足条件，N代表N日前满足条件)
def BARSNEXT(s:Union[List, Sequence, np.ndarray, pd.Series], asArray:bool = False):
    """
    BARSNEXT: 未来多少周期才能再次满足条件
    -----------------------------
    Params:
        s - 条件判断序列
        asArray - 是否转为np.array
    
    Return:
        未来满足条件离当前k周期数
    """

    r = _asSeries(s)[::-1]
    i = len(r)-r.index.argsort()
    p = (i * r).cummax()
    return _ensure(((p>0) * (i-p) - (p<1))[::-1], asArray)

# 判断两条线是否相交
def CROSS(a:Union[List, Sequence, np.ndarray, pd.Series], b:Union[List, Sequence, np.ndarray, pd.Series], asArray:bool = False):               
    """
    CROSS: 均线交叉
    -----------------------------
    Params:
        a - 均线A
        b - 均线B
        asArray - 是否转为np.array
    
    Return:
        a > b and ref(a,1)<=ref(b,1)
    """

    r = (_asSeries(a) > _asSeries(b)).astype(int)
    return r.diff()==1

#####################################################################
## 6.技术指标
#####################################################################

def MACD(s:Union[List, Sequence, np.ndarray, pd.Series], short:int = 12, long:int = 26, mid:int = 9, asArray:bool = False):
    """
    MACD: 指数平滑异同移动平均线
    -------------------------------------
    DIF:EMA(CLOSE,SHORT)-EMA(CLOSE,LONG);
    DEA:EMA(DIF,MID);
    MACD:(DIF-DEA)*2,COLORSTICK;
    """

    r = _ensure(s, asArray)
    dif = EMA(r, short) - EMA(r, long)
    dea = EMA(dif, mid)
    macd = (dif - dea)*2
    return dif, dea, macd    

def BOLL(s:Union[List, Sequence, np.ndarray, pd.Series], m:int = 20, asArray:bool = False):
    """
    BOLL: 布林线
    --------------------------------------
    BOLL:MA(CLOSE,M);
    UB:BOLL+2*STD(CLOSE,M);
    LB:BOLL-2*STD(CLOSE,M);
    """
    r = _ensure(s, asArray)
    boll = MA(s, m)
    std = STD(s, m)
    ub = boll + 2 * std
    lb = boll - 2 * std 
    return boll, ub, lb        

def KDJ(bars:pd.DataFrame, n:int = 9, m1:int = 3, m2:int = 3, asArray:bool = False):
    """
    KDJ: 随机指标
    --------------------------------------
    RSV:=(CLOSE-LLV(LOW,N))/(HHV(HIGH,N)-LLV(LOW,N))*100;
    K:SMA(RSV,M1,1);
    D:SMA(K,M2,1);
    J:3*K-2*D;
    """
    close = CLOSE(bars, asArray)
    high = HIGH(bars, asArray)
    high_hhv = HHV(high, n, asArray)
    low = LOW(bars, asArray)
    low_llv = LLV(low, n, asArray)
    rsv = (close - low_llv)/(high_hhv-low_llv) * 100
    k = SMA(rsv, m1, 1, asArray)
    d = SMA(k, m2, 1, asArray)
    j = 3*k - 2*d
    return k, d, j            

def MTM(s:Union[List, Sequence, np.ndarray, pd.Series], n:int = 12, m:int = 6, asArray:bool = False):
    """
    MTM: 动量线
    --------------------------------------
    MTM:CLOSE-REF(CLOSE,MIN(BARSCOUNT(C),N));
    MTMMA:MA(MTM,M);
    """
    close = _ensure(s, asArray)
    mtm = REF(close, np.minimum(BARSCOUNT(close),n))
    mtmma = MA(mtm,m)
    return mtm, mtmma          


def OBV(bars:pd.DataFrame, m:int = 30, asArray:bool = False):
    """
    OBV: 能量潮（累积能量线）
    --------------------------------------
    VA:=IF(CLOSE>REF(CLOSE,1),VOL,-VOL);
    OBV:SUM(IF(CLOSE=REF(CLOSE,1),0,VA),0);
    MAOBV:MA(OBV,M);
    """
    close = CLOSE(bars, asArray)
    ref_close = REF(close, 1)
    vol = VOLUME(bars, asArray)
    va = IF(close > ref_close, vol, -vol)
    obv = np.cumsum((close!=ref_close)*va)
    maobv = MA(obv, m)
    return obv, maobv